# 导入LangChain核心消息类型，用于构建对话消息
from langchain_core.messages import AIMessage, SystemMessage,HumanMessage, ToolMessage
# 导入LangGraph的Send类型，用于消息传递
from langgraph.types import Send
# 导入LangGraph的状态图相关类，用于构建工作流图
from langgraph.graph import StateGraph, MessagesState
# 导入图的起始和结束节点常量
from langgraph.graph import START, END
# 导入消息合并函数，用于自动处理消息列表的合并
from langgraph.graph.message import add_messages
# 导入自定义的大语言模型配置函数
from llm import Tongyi
# 导入系统提示词常量
from prompts import SYSTEM_PROMPT
# 导入类型注解相关工具
from typing import Annotated, TypedDict
# 导入工具函数列表
from tools import tools

# 加载大模型实例
llm_Tongyi = Tongyi()

# 定义状态结构
class OverallState(TypedDict):
    # 消息历史
    messages: Annotated[list, add_messages] # 消息历史列表，自动合并
    # 待执行动作
    actions: list[dict]
    # 最终输出 
    output: str
    # 用户输入
    user_prompt: str

# 定义codeact模式主类
class CodeActGraph():
    # 运行主方法
    def run(self, user_prompt: str):
        def _create_agent_action(ai_message):
            # 提取代码部分
            content = ai_message.content
            if "```python" in content:
                code_blocks = content.split("```python")
                code = code_blocks[1].split("```")[0].strip()
                return {"tool": "execute_python", "tool_input": code}
            return None

        # 定义llm_call节点
        def llm_call(state: OverallState):
            # 构建包含系统提示词和用户提示词的消息列表
            messages = [
                SystemMessage(content=SYSTEM_PROMPT),  # 创建系统消息，设置AI的角色和指令
                HumanMessage(content=state["user_prompt"]) # 创建用户消息，包含用户的问题
            ]+state["messages"] # 添加历史消息

            ret = llm_Tongyi.invoke(messages)  # 调用大模型生成回复

            action = _create_agent_action(ret)
            if action:
                return {"messages": [ret], "actions": [action], "output": ""}
            return {"messages": [ret], "output": ret.content}

        # 定义enter_process节点
        def enter_process(state: OverallState):
            # 判断是否需要执行代码
            if state["output"] != "":
                return END
            return "process_node"

        # 定义过程调用tool节点
        def process_node(state: OverallState):
            # 遍历待执行的动作列表
            actions = state["actions"]
            for action in actions:
                tool_name = action["tool"]
                tool_input = action["tool_input"]
                # 根据工具名称找到对应的工具函数
                tool_fn = next(t for t in tools if t.name == tool_name)
                observation = tool_fn.invoke(tool_input)    #执行工具
                state["messages"].append(HumanMessage(content=f"##执行结果:\n{observation}")) # 添加执行结果到消息历史

            return state

        # 实例化状态图
        graph = StateGraph(OverallState)
        # 3个节点：llm_call（LLM调用）、enter_process（流程判断）、process_node（代码执行）
        graph.add_node("llm_call", llm_call) # 调用大模型
        graph.add_node("enter_process", enter_process) # 判断是否需要执行代码
        graph.add_node("process_node", process_node) # 执行代码

        graph.add_edge(START, "llm_call") # 设置入口节点
        graph.add_conditional_edges("llm_call",enter_process) # 判断是否需要执行代码
        graph.add_edge("process_node", "llm_call") # 执行代码

        agent = graph.compile() # 编译图
        ret = agent.invoke(OverallState(user_prompt=user_prompt))
        
        return ret["output"]

if __name__ == "__main__":
    graph = CodeActGraph()
    ret = graph.run("请计算斐波那契数列中第5位的数是多少，第1位是1")
    
    # tips：绘图最好在本机上展示，不要在远端服务器
    #ret = graph.run("某上市公司的2022年毛利率为66.75%，2023年毛利率为44.07%，2024年毛利率为42.94%，请完成趋势折线图的绘制")
    print(ret)
